Add reflect and symmetric padding modes to mx.pad by katlun-lgtm · Pull Request #3608 · ml-explore/mlx

katlun-lgtm · 2026-05-30T20:13:57Z

Summary

Adds numpy.pad-compatible \"reflect\" and \"symmetric\" modes to mx.pad. These join the existing \"constant\" and \"edge\" modes.

Both modes match numpy.pad semantics for arbitrary pad sizes — when the pad width exceeds the axis length the reflection repeats, exactly as NumPy does. (Earlier attempts at these modes were limited to pad < dim; this implementation removes that restriction.)

reflect — mirror padding that does not repeat the edge value (period 2(n-1)).
symmetric — mirror padding that does repeat the edge value (period 2n).

Motivation

This was developed as part of porting mobileportrait-mlx (a TPS-based face-animation model) to pure MLX for training and inference on Apple Silicon. The full pipeline — CelebV-HQ frame extraction, keypoint cache, TPS warping, perceptual loss — trains end-to-end in MLX on an M3 Max. After training on 3,000 clips the model runs inference at 22.6 fps (256×256, dense-motion + inpainting stage, M3 Max) and reaches render L1 = 0.0566 on held-out clips.

Implementation

reflect_pad (in mlx/ops.cpp) builds a per-axis index map with a triangle-wave reflection function and gathers with take — one take per padded axis. A degenerate axis of length 1 maps every coordinate to 0. No new primitive or kernel is introduced; it composes existing ops, so it works on every backend and is differentiable for free.

Files changed

mlx/ops.cpp — reflect_pad helper + reflect / symmetric dispatch branches in pad.
python/src/ops.cpp — extend the mode Literal and docstring.
python/tests/test_ops.py — test_pad_reflect_symmetric.
tests/ops_tests.cpp — reflect/symmetric CHECK cases incl. multi-reflect.

Testing

All run locally on an M3 Max:

Python test_pad_reflect_symmetric — 13 shape/pad-width cases × 2 modes compared element-for-element against numpy.pad (in-bounds, multi-reflect where pad ≫ axis, asymmetric per-axis, zero-width sides, and degenerate axes n==1, n==2). Exact match.
C++ tests/ops_tests.cpp "test pad" — 9 assertions pass.
Full C++ suite — 251 cases / 251 passed, 3442 assertions / 0 failed.

>>> import mlx.core as mx
>>> a = mx.array([1, 2, 3])
>>> mx.pad(a, 2, mode=\"reflect\")
array([3, 2, 1, 2, 3, 2, 1], dtype=int32)
>>> mx.pad(a, 2, mode=\"symmetric\")
array([2, 1, 1, 2, 3, 3, 2], dtype=int32)

zcbenz · 2026-06-01T06:33:33Z

Thanks for the PR, there was actually #2862 which is still waiting for a thorough review to merge. I think your PR is more elegant and I'm more in favor of this one, just in case do you think we are missing anything from #2862?

katlun-lgtm · 2026-06-13T03:01:20Z

Thanks @zcbenz! I went through #2862 carefully — functionally the two cover the same ground (reflect + symmetric, matching numpy.pad including the edge-excluded vs edge-included distinction), so nothing user-facing is missing.

Differences I found:

Implementation: [Feature] Additional padding modes #2862 builds the padding with a slice/concatenate/tile/slice_update chain (a separate function per mode). This PR instead computes a single triangle-wave index map per padded axis and does one take (gather), with both modes sharing one helper via an include_edge flag — fewer lines and a lighter op graph.
Multi-reflect (pad wider than the axis): handled in both; here it falls out of the modular index map rather than needing explicit tiling/reps logic.
Tests: the suite here is a superset — it adds degenerate axes (n == 1, n == 2), multi-reflect on both sides simultaneously, asymmetric per-axis widths, and 3D cases, all checked elementwise against numpy.pad.
The one thing [Feature] Additional padding modes #2862 had that this PR didn't — an ACKNOWLEDGMENTS.md entry — I've just added.

Happy to adjust naming or docs to whatever you prefer.

zcbenz

Thanks!

katlun-lgtm · 2026-07-07T02:58:52Z

Following up now that #2862 has been closed — that leaves this PR as the remaining path for reflect / symmetric padding modes. It's approved with all CI green; happy to rebase onto latest main if you'd like it current, otherwise it's ready to merge whenever convenient. Thanks!

katlun-lgtm · 2026-07-31T02:07:13Z

Rebased onto current main — this had drifted into a conflict, so it wasn't mergeable regardless of review. It's clean now at f38f4522.

The only conflict was documentation: main renamed the documented argument constant_value to constant_values to match the actual keyword, while this branch still had the old spelling. I resolved in main's favour and kept the two new mode lines above it. The implementation, the C++ tests and the Python tests all applied unchanged, and nothing in mlx/ops.cpp has touched pad since this branch was cut.

For the record on the earlier hold: #2862 was closed unmerged on July 1, so this is the remaining path for reflect / symmetric. The June 13 approval carried across the rebase.

The new head has no checks on it yet — fork runs need an approval to start, if someone can grant it.

Implements numpy.pad-compatible "reflect" and "symmetric" modes for mx.pad, matching numpy semantics for arbitrary pad sizes (the reflection repeats when the pad width exceeds the axis length). - mlx/ops.cpp: reflect_pad helper builds a per-axis triangle-wave index map and gathers with take; one take per padded axis. reflect uses period 2(n-1) and skips the edge; symmetric uses period 2n and repeats the edge. n==1 maps to 0. - python/src/ops.cpp: extend the pad mode Literal and docstring. - python/tests/test_ops.py: test_pad_reflect_symmetric covers in-bounds, multi-reflect, asymmetric per-axis, zero-width sides, and degenerate axes (n==1, n==2), checked against numpy.pad. - tests/ops_tests.cpp: reflect/symmetric CHECK cases incl. multi-reflect.

zcbenz

Sorry for the late response, I had to find someone to verify the correctness of the implementation before merging, since no one did I have to spend some time myself.

So this is not a performant implementation: it builds the indices with the same length of output in the host, which would be rather slow for large inputs.

A performant implementation should follow edge_pad: use slice_update to write paddings and unpadded array to output, with host calculating only the start/stop indices. The reflect_pad only needs to change how the paddings are generated, by slicing and then flipping, you can find some idea in JAX's implementation https://github.com/jax-ml/jax/blob/8d1be7d7d1c4bd19a0de12c917709dd20f2de34d/jax/_src/numpy/lax_numpy.py#L3988.

Follows edge_pad's pattern: place the input into a zero-filled output with slice_update, then extend each axis outward via slice+flip instead of building a host-side index array of output length and gathering. A pad width larger than the axis loops in tiles, re-slicing the data just written to continue the periodic reflect/symmetric pattern (matches numpy.pad, including multi-reflect). Addresses zcbenz's review on ml-explore#3608: sub-ms on 5M/4M-element arrays for the common in-bounds case, vs. an O(output size) host loop + gather before. All 154 test_ops.py cases pass, including the existing multi-reflect coverage that exercises the new tiling loop.

katlun-lgtm · 2026-08-11T12:15:54Z

Rewrote it to follow edge_pad's pattern, pushed as beec433e.

Same structure as edge_pad: place the input into a zero-filled output via slice_update, then extend each axis outward with slice + flip on already-written data. Host only computes start/stop indices, no per-element index array. For a pad width larger than the axis, it loops in tiles, re-slicing the data it just wrote to continue the periodic pattern — same idea as your JAX reference, just landing in a fixed slice_update buffer instead of growing via concatenate.

All 154 test_ops.py cases pass, including the existing multi-reflect coverage (pad > axis) that exercises the tiling loop. Benchmarked on M3 Max: 5M-element 1D and 2000x2000 2D arrays with small in-bounds padding are sub-millisecond now (0.7-1.0ms); a deliberately pathological multi-reflect case (5M elements, pad 2M each side, 3 tiles) is still ~1.1ms.

CI's clang-format hook flagged the two dispatcher lines as too long.

katlun-lgtm · 2026-08-11T12:35:05Z

Fixed the clang-format failure from the last run — my two reflect_pad(...) call sites in pad() were too long, wrapped them to match what the hook wanted. Pushed as 0f87956b. Also picked up the main-merge that landed on this branch in the meantime.

katlun-lgtm marked this pull request as ready for review May 30, 2026 20:22

zcbenz mentioned this pull request Jun 10, 2026

Add stft and istft to mlx.core.fft #3639

Open

4 tasks

zcbenz force-pushed the add-reflect-symmetric-pad branch from 63a3e46 to 579bfd3 Compare June 13, 2026 00:11

zcbenz approved these changes Jun 13, 2026

View reviewed changes

zcbenz mentioned this pull request Jul 1, 2026

[Feature] Additional padding modes #2862

Closed

4 tasks

katlun-lgtm force-pushed the add-reflect-symmetric-pad branch from 71a68c6 to f38f452 Compare July 31, 2026 01:35

zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 4, 2026

katlun-lgtm and others added 3 commits August 11, 2026 00:25

Fix lint

639b706

docs: add ACKNOWLEDGMENTS entry for reflect/symmetric pad modes

f2ef917

katlun-lgtm force-pushed the add-reflect-symmetric-pad branch from f38f452 to f2ef917 Compare August 11, 2026 04:25

zcbenz requested changes Aug 11, 2026

View reviewed changes

zcbenz removed the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 11, 2026

katlun-lgtm and others added 2 commits August 11, 2026 08:22

Merge branch 'main' into add-reflect-symmetric-pad

6a624b4

style: wrap reflect_pad call sites to satisfy clang-format

0f87956

CI's clang-format hook flagged the two dispatcher lines as too long.

katlun-lgtm requested a review from zcbenz August 11, 2026 12:38

nit

1e70367

zcbenz approved these changes Aug 11, 2026

View reviewed changes

zcbenz merged commit 9f35f77 into ml-explore:main Aug 11, 2026

zcbenz mentioned this pull request Aug 11, 2026

[Feature Request] Add Symmetric and Reflect padding modes to mx.pad #2155

Closed

katlun-lgtm deleted the add-reflect-symmetric-pad branch August 11, 2026 23:51

codeAnqiang-ma mentioned this pull request Aug 13, 2026

[BUG] mx.pad with mode="reflect"/"symmetric" hangs in an infinite loop when a padded axis is empty #4221

Closed

Navigation Menu

Uh oh!

Add reflect and symmetric padding modes to mx.pad - #3608

Add reflect and symmetric padding modes to mx.pad#3608
zcbenz merged 7 commits into
ml-explore:mainfrom
katlun-lgtm:add-reflect-symmetric-pad

katlun-lgtm commented May 30, 2026 •

edited

Loading

Uh oh!

zcbenz commented Jun 1, 2026

Uh oh!

katlun-lgtm commented Jun 13, 2026

Uh oh!

zcbenz left a comment

Uh oh!

katlun-lgtm commented Jul 7, 2026 •

edited

Loading

Uh oh!

katlun-lgtm commented Jul 31, 2026

Uh oh!

zcbenz left a comment

Uh oh!

katlun-lgtm commented Aug 11, 2026

Uh oh!

katlun-lgtm commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

katlun-lgtm commented May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Implementation

Files changed

Testing

Uh oh!

zcbenz commented Jun 1, 2026

Uh oh!

katlun-lgtm commented Jun 13, 2026

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

katlun-lgtm commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

katlun-lgtm commented Jul 31, 2026

Uh oh!

zcbenz left a comment

Choose a reason for hiding this comment

Uh oh!

katlun-lgtm commented Aug 11, 2026

Uh oh!

katlun-lgtm commented Aug 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

katlun-lgtm commented May 30, 2026 •

edited

Loading

katlun-lgtm commented Jul 7, 2026 •

edited

Loading